##########################################################################################

library(data.table)
library(optparse)
library(Seurat)
library(ArchR)
library(BSgenome.Hsapiens.UCSC.hg38)
library(org.Hs.eg.db)
library(GO.db)
library(ggthemes)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--comine_data_all_file"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/subcell/cluster5/cluster5.combineRNA.motif_peak2gene.Rdata"

    ## 所有的细胞的,计算maxdelt
    comine_data_all_file <- "~/20231121_singleMuti/results/qc_atac/testis_combined_peak.combineRNA.qc.Rdata"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/tf_regulators/cluster5"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
comine_data_all_file <- opt$comine_data_all_file
scriptPath <- opt$scriptPath
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################
## 导入数据
a <- load(comine_data_file)
## atac_proj

b <- load(comine_data_all_file)
## testis_combined_peak_combineRNA

## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

##########################################################################################
# Identify regulatory targets of TFs 
##########################################################################################
# ChromVAR deviations matrix: (rows motif names x cols cell names)
motifMatrix <- getMatrixFromProject(atac_proj, useMatrix="MotifMatrix")
## motfi在每个细胞中的活性程度
deviationsMatrix <- assays(motifMatrix)$deviations

# GeneIntegration Matrix: (rows gene names x cols cell names)
GIMatrix <- getMatrixFromProject(atac_proj, useMatrix="GeneExpressionMatrix")
GImat <- assays(GIMatrix)$GeneExpressionMatrix
rownames(GImat) <- rowData(GIMatrix)$name
# Remove unexpressed genes
GImat <- as(GImat[Matrix::rowSums(GImat) > 0,], "sparseMatrix") 

# Identify pseudobulks for performing matrix correlations
knn_groups <- getLowOverlapAggregates(atac_proj, target.agg=500, k=100, overlapCutoff=0.8)

kgrps <- unique(knn_groups$group)

# GeneIntegrationMatrix
GIMatPsB <- lapply(kgrps, function(x){
  use_cells <- knn_groups[knn_groups$group==x,]$cell_name
  Matrix::rowMeans(GImat[,use_cells])
  }) %>% do.call(cbind,.)
colnames(GIMatPsB) <- kgrps

# In rare instances, we can get pseudo-bulked genes that have zero averages
GIMatPsB <- GIMatPsB[Matrix::rowSums(GIMatPsB) > 0,]

# DeviationsMatrix
DevMatPsB <- lapply(kgrps, function(x){
  use_cells <- knn_groups[knn_groups$group==x,]$cell_name
  Matrix::rowMeans(deviationsMatrix[,use_cells])
  }) %>% do.call(cbind,.)
colnames(DevMatPsB) <- kgrps

# Perform chromVAR deviations to Integrated RNA correlation analysis:
start <- Sys.time()
geneCorMat <- cor2Matrices(DevMatPsB, GIMatPsB)
colnames(geneCorMat) <- c("motifName", "symbol", "Correlation", "FDR")
end <- Sys.time()
message(sprintf("Finished correlations in %s minutes.", round((end  - start)/60.0, 2)))

# Already filtered to only expressed genes
allGenes <- rownames(GIMatPsB) %>% sort() 

# Get locations of motifs of interest:
motifPositions <- getPositions(atac_proj, name="Motif")
motifGR <- stack(motifPositions, index.var="motifName")

# Get peak to gene GR
corrCutoff <- 0.45 # Used in labeling peak2gene links
p2gGR <- getP2G_GR(atac_proj, corrCutoff=corrCutoff)

## function
calculateLinkageScore <- function(motifLocs, p2gGR){
  # Calculate Linkage Score (LS) for each gene in p2gGR with regards to a motif location GR
  ###################################
  # For a given gene, the LS = sum(corr peak R2 * motifScore)
  ol <- findOverlaps(motifLocs, p2gGR, maxgap=0, type=c("any"), ignore.strand=TRUE)
  olGenes <- p2gGR[to(ol)]
  olGenes$motifScore <- motifLocs[from(ol)]$score
  olGenes$R2 <- olGenes$Correlation**2 # All p2g links here are already filtered to only be positively correlated
  LSdf <- mcols(olGenes) %>% as.data.frame() %>% group_by(symbol) %>% summarise(LS=sum(R2*motifScore)) %>% as.data.frame()
  LSdf <- LSdf[order(LSdf$LS, decreasing=TRUE),]
  LSdf$rank <- 1:nrow(LSdf)
  return(LSdf)
}

calculateMotifEnrichment <- function(motifLocs, p2gGR){
  # Calculate Motif enrichment per gene
  ###################################
  # For a given gene, calculate the hypergeometric enrichment of motifs in 
  # linked peaks (generally will be underpowered)
  motifP2G <- p2gGR[overlapsAny(p2gGR, motifLocs, maxgap=0, type=c("any"), ignore.strand=TRUE)]
  m <- length(motifP2G) # Number of possible successes in background
  n <- length(p2gGR) - m # Number of non-successes in background

  motifLinks <- motifP2G$symbol %>% getFreqs()
  allLinks <- p2gGR$symbol %>% getFreqs()
  df <- data.frame(allLinks, motifLinks=motifLinks[names(allLinks)])
  df$motifLinks[is.na(df$motifLinks)] <- 0
  df$mLog10pval <- apply(df, 1, function(x) -phyper(x[2]-1, m, n, x[1], lower.tail=FALSE, log.p=TRUE)/log(10))
  df <- df[order(df$mLog10pval, decreasing=TRUE),]
  df$symbol <- rownames(df)
  return(df)
}

markerGenes <- c("UTF1", "KIT", "STRA8", 
"SPO11", "SYCP3", "OVOL2", 
"NME8" , "TXNDC8" ,"TNP1" , 
"PRM1" ,"AMH", "DLK1",
"MYH11","NOTCH3","CD14",
"VWF","NKG7" ,"FGFBP2")

#精原细胞(UTF1) —Cell Research. 2018
#正在分化的精原细胞(KIT)—Cell stem cell. 2018
#分化完成的精原细胞(STRA8)—Cell stem cell. 2018
#细线期精母细胞(SPO11)—Cell stem cell. 2018
#偶线期精母细胞(SYCP3)—Cell Reports. 2018
#粗线期精母细胞(OVOL2)—Cell stem cell. 2018
#双线期精母细胞(NME8)—Cell stem cell. 2018
#圆形精子和长形精子细胞(TXNDC8)    —Cell stem cell. 2018
#精子(TNP1、PRM1)—Cell stem cell. 2018
#Sertoli细胞(AMH)—Cell stem cell. 2018
#Leydig细胞(DLK1)—Cell Research. 2018
#肌样细胞(MYH11)—Cell stem cell. 2018
#周细胞(NOTCH3)—Human Molecular Genetics.2022
#巨噬细胞(CD14)—Cell Research. 2018
#内皮细胞(VWF)—Cell Research. 2018
#NKT细胞(NKG7、FGFBP2)

##########################################################################################
# plot all TF regulators
regPlotDir <- paste0(out_path, "/TFregulatorPlots")
dir.create(regPlotDir, showWarnings = FALSE, recursive = TRUE)

# Store results for each TF
res_list <- list()

## 展示所有的
regulators <- levels(motifGR$motifName)

for(motif in regulators){
  motif_short <- strsplit(motif,"_")[[1]][1]
  # First get motif positions
  motifLocs <- motifGR[motifGR$motifName == motif]
  # Calculate Linkage Score for motif
  LS <- calculateLinkageScore(motifLocs, p2gGR)
  # Get just genes correlated to motif
  motifGeneCorDF <- geneCorMat[geneCorMat$motifName == motif,]
  plot_df <- merge(LS, motifGeneCorDF, by="symbol", all.x=TRUE)
  # Calculate motif enrichment per gene
  ME <- calculateMotifEnrichment(motifLocs, p2gGR)
  plot_df <- merge(plot_df, ME, by="symbol", all.x=TRUE)
  plot_df2 <- plot_df

  plot_df <- plot_df[,c("symbol", "LS", "Correlation", "FDR", "mLog10pval")]
  plot_df$toLabel <- "NO"
  topN <- 5
  plot_df <- plot_df[order(plot_df$LS, decreasing=TRUE),]
  plot_df$rank_LS <- 1:nrow(plot_df)
  plot_df$toLabel[1:topN] <- "YES"
  plot_df <- plot_df[order(plot_df$Correlation, decreasing=TRUE),]
  plot_df$rank_Corr <- 1:nrow(plot_df)
  plot_df$toLabel[1:topN] <- "YES"
  plot_df <- plot_df[order(plot_df$mLog10pval, decreasing=TRUE),]
  plot_df$rank_Pval <- 1:nrow(plot_df)
  plot_df$toLabel[1:10] <- "YES"
  plot_df$meanRank <- apply(plot_df[,c("rank_LS", "rank_Corr", "rank_Pval")], 1, mean)
  plot_df <- plot_df[order(plot_df$meanRank, decreasing=FALSE),]
  plot_df$toLabel[1:topN] <- "YES"
  # Label any marker genes in window of interest
  LS_window <- quantile(plot_df$LS, 0.8)
  corr_window <- 0.25
  pos_top_genes <- plot_df[plot_df$LS > LS_window & plot_df$Correlation > corr_window,]$symbol
  neg_top_genes <- plot_df[plot_df$LS > LS_window & -plot_df$Correlation > corr_window,]$symbol
  if(nrow(plot_df[plot_df$symbol %in% c(pos_top_genes, neg_top_genes) & plot_df$symbol %in% markerGenes,]) > 0){
    plot_df[plot_df$symbol %in% c(pos_top_genes, neg_top_genes) & plot_df$symbol %in% markerGenes,]$toLabel <- "YES"
  }
  res_list[[motif_short]] <- pos_top_genes # Save regulatory targets
  # Save dataframe of results
  save_df <- plot_df[plot_df$symbol %in% c(pos_top_genes, neg_top_genes),c(1:5)]
  save_df <- save_df[order(save_df$Correlation, decreasing=TRUE),]
  saveRDS(save_df, paste0(regPlotDir, sprintf("/regulatory_targets_%s.rds", motif_short)))
  plot_df <- plot_df[order(plot_df$mLog10pval, decreasing=FALSE),]
  # Label motif as well
  plot_df$toLabel[which(plot_df$symbol == motif_short)] <- "YES"
  plot_df$symbol[which(plot_df$toLabel == "NO")] <- ""
  # Threshold pvalue for plotting
  maxPval <- 5
  plot_df$mLog10pval <- ifelse(plot_df$mLog10pval > maxPval, maxPval, plot_df$mLog10pval)
  #Plot results
  p <- (
    ggplot(plot_df, aes(x=Correlation, y=LS, color=mLog10pval)) 
      #+ geom_point(size = 2)
      + ggrastr::geom_point_rast(size=2)
      + ggrepel::geom_text_repel(
          data=plot_df[plot_df$toLabel=="YES",], aes(x=Correlation, y=LS, label=symbol), 
          #data = plot_df, aes(x=Correlation, y=LS, label=symbol), #(don't do this, or the file will still be huge...)
          size=2,
          point.padding=0, # additional pading around each point
          box.padding=0.5,
          min.segment.length=0, # draw all line segments
          max.overlaps=Inf, # draw all labels
          #nudge_x = 2,
          color="black"
      ) 
      + geom_vline(xintercept=0, lty="dashed") 
      + geom_vline(xintercept=corr_window, lty="dashed", color="red")
      + geom_vline(xintercept=-corr_window, lty="dashed", color="red")
      + geom_hline(yintercept=LS_window, lty="dashed", color="red")
      + theme_BOR(border=FALSE)
      + theme(panel.grid.major=element_blank(), 
              panel.grid.minor= element_blank(), 
              plot.margin = unit(c(0.25,1,0.25,1), "cm"), 
              aspect.ratio=1.0,
              #legend.position = "none", # Remove legend
              axis.text.x = element_text(angle=90, hjust=1))
      + ylab("Linkage Score") 
      + xlab("Motif Correlation to Gene") 
      + scale_color_gradientn(colors=cmaps_BOR$zissou, limits=c(0, maxPval))
      + scale_y_continuous(expand = expansion(mult=c(0,0.05)))
      + scale_x_continuous(limits = c(-0.85, 0.955)) # Force plot limits
      + ggtitle(sprintf("%s putative targets", motif_short))
      )

  # Positively regulated genes:
  if( length(pos_top_genes) > 0 ){
    upGO <- rbind(
      calcTopGo(allGenes, interestingGenes=pos_top_genes, nodeSize=5, ontology="BP"), 
      calcTopGo(allGenes, interestingGenes=pos_top_genes, nodeSize=5, ontology="MF")
      )
    if(min(upGO$pvalue) < 1){
      upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]
      up_go_plot <- topGObarPlot(upGO, cmap=cmaps_BOR$comet, nterms=6, border_color="black", 
        barwidth=0.9, title=sprintf("%s putative targets (%s genes)", motif_short, length(pos_top_genes)), enrichLimits=c(0, 6))
    }else{
      up_go_plot <- ggplot()
    }
  }

  # Negatively regulated genes:
  if( length(neg_top_genes) > 0 ){
    downGO <- rbind(
      calcTopGo(allGenes, interestingGenes=neg_top_genes, nodeSize=5, ontology="BP"), 
      calcTopGo(allGenes, interestingGenes=neg_top_genes, nodeSize=5, ontology="MF")
      )

    ## 不存在p值显著的通路则不展示
    if(min(downGO$pvalue) < 1){
      downGO <- downGO[order(as.numeric(downGO$pvalue), decreasing=FALSE),]
      down_go_plot <- topGObarPlot(downGO, cmap=cmaps_BOR$comet, nterms=6, border_color="black", 
        barwidth=0.9, title=sprintf("%s putative targets (%s genes)", motif_short, length(neg_top_genes)), enrichLimits=c(0, 6))
    }else{
      down_go_plot <- ggplot()
    }
  }

  pdf(paste0(regPlotDir, sprintf("/%s_LS.pdf", motif_short)), width=8, height=6)
  print(p)
  if( length(pos_top_genes) > 0 ){
    print(up_go_plot)
  }
  if( length(neg_top_genes) > 0 ){
    print(down_go_plot)
  }
  dev.off()

  ## 标记TF的靶基因
  plot_df2$label_target <- "NO"
  plot_df2$putative_targets <- "NO"
  ## LS最高的前5个、相关系数最高的前5个、motif富集程度p最显著的前10个、综合LS和相关系数以及富集的均值前10个
  plot_df2[plot_df2$symbol %in% plot_df$symbol , "label_target"] <- "YES"
  ## LS > 80%分位数且关联系数>0.25或者<-0.25
  plot_df2[plot_df2$symbol %in% c(pos_top_genes , neg_top_genes) , "putative_targets"] <- "YES"

  out_file <- paste0(regPlotDir, sprintf("/%s_LS.tsv", motif_short))
  write.table(plot_df2 , out_file , row.names = F , sep = "\t" , quote = F)
}